import os
import base64
from PIL import Image
import io
import logging
from config import SUPPORTED_IMAGE_FORMATS, MAX_IMAGE_SIZE

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class MultimodalProcessor:
    def __init__(self):
        """初始化多模态处理器"""
        pass

    def is_image_file(self, file_path):
        """检查文件是否为支持的图像格式"""
        ext = os.path.splitext(file_path)[1].lower()
        return ext in SUPPORTED_IMAGE_FORMATS

    def process_image(self, image_path, max_size=None):
        """
        处理图像：调整大小并转换为Base64编码

        参数:
            image_path: 图像文件路径
            max_size: 最大尺寸元组 (width, height)，默认为配置中的值

        返回:
            处理后的图像Base64编码字符串，或None（处理失败）
        """
        max_size = max_size or MAX_IMAGE_SIZE

        try:
            with Image.open(image_path) as img:
                # 转换为RGB模式（处理透明背景）
                if img.mode in ('RGBA', 'LA'):
                    background = Image.new(img.mode[:-1], img.size, (255, 255, 255))
                    background.paste(img, img.split()[-1])
                    img = background
                elif img.mode == 'P':
                    img = img.convert('RGB')

                # 调整大小
                img.thumbnail(max_size)

                # 转换为Base64
                buffered = io.BytesIO()
                img.save(buffered, format="JPEG")
                img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')

                logger.info(f"成功处理图像: {image_path}, 尺寸: {img.size}")
                return img_base64
        except Exception as e:
            logger.error(f"处理图像 {image_path} 失败: {str(e)}")
            return None

    def create_multimodal_prompt(self, text_content, image_paths=None):
        """
        创建多模态提示词，包含文本和图像信息

        参数:
            text_content: 文本内容
            image_paths: 图像文件路径列表

        返回:
            多模态提示词内容
        """
        prompt_parts = [text_content]

        if image_paths and isinstance(image_paths, list):
            for i, image_path in enumerate(image_paths):
                if self.is_image_file(image_path):
                    img_base64 = self.process_image(image_path)
                    if img_base64:
                        # 在实际应用中，这里的格式取决于API要求
                        # 这里使用一种假设的格式表示图像
                        prompt_parts.append(f"\n[图像 {i + 1}]: {img_base64[:50]}...")  # 只显示部分Base64作为示意

        return "\n".join(prompt_parts)

    def process_input(self, input_data):
        """
        处理输入数据，判断类型并进行相应处理

        参数:
            input_data: 输入数据，可以是文本、文件路径或列表

        返回:
            处理后的内容
        """
        if isinstance(input_data, list):
            # 处理多个输入
            text_parts = []
            image_paths = []

            for item in input_data:
                if isinstance(item, str):
                    if os.path.exists(item) and self.is_image_file(item):
                        image_paths.append(item)
                    else:
                        text_parts.append(item)

            text_content = "\n".join(text_parts)
            return self.create_multimodal_prompt(text_content, image_paths)

        elif isinstance(input_data, str):
            if os.path.exists(input_data) and self.is_image_file(input_data):
                # 单一图像输入
                return self.create_multimodal_prompt("分析以下图像:", [input_data])
            else:
                # 文本输入
                return input_data

        else:
            logger.error(f"不支持的输入类型: {type(input_data)}")
            return None
